#include <sstream>
using namespace std;

#include "ndim-utils.h"
#include <misc/error.h>

NDimRangeTraverseCB::~NDimRangeTraverseCB() 
{
}

NDimRange::NDimRange(const vector<size_t>& region_shape,
		     const vector<size_t>& shape) 
{
  this->recurse_level = 0;
  this->region_shape.resize(region_shape.size());
  for (size_t i_dim = 0; i_dim < region_shape.size(); i_dim++) {
    this->region_shape[i_dim] = make_pair(0, region_shape[i_dim]);
  }
  this->shape = shape;
  this->curr_coor.resize(shape.size(), 0);
}

NDimRange::NDimRange(const vector<pair<size_t, size_t> >& region_shape,
		     const vector<size_t>& shape) 
{
  this->recurse_level = 0;
  this->region_shape = region_shape;
  this->shape = shape;
  this->curr_coor.resize(shape.size(), 0);
}

void NDimRange::traverse(NDimRangeTraverseCB* cb) 
{
  fill(curr_coor.begin(), curr_coor.end(), 0);
  this->traverse_dimension(cb, 0);
}

void NDimRange::traverse_dimension(NDimRangeTraverseCB* cb, size_t i_dim) 
{
  recurse_level++;
  for (size_t i = region_shape[i_dim].first;
       i < min(region_shape[i_dim].second, shape[i_dim]);
       i++) {
    curr_coor[i_dim] = i;
    if (recurse_level == shape.size()) {
      cb->do_it(curr_coor);
    }
    if (i_dim + 1 != shape.size()) {
      this->traverse_dimension(cb, i_dim + 1);
    }
  }
  recurse_level--;
}

static size_t get_direct_index__(size_t tot, 
				 const vector<size_t>& shape,
				 const vector<size_t>& coor,
				 size_t i_dim)
{
  if (i_dim + 1 != shape.size()) {
    tot += coor[i_dim];
    tot *= shape[i_dim + 1];
    return get_direct_index__(tot, shape, coor, i_dim + 1);
  }
  tot += coor[i_dim];
  return tot;
}

size_t get_direct_index(const vector<size_t>& shape,
			const vector<size_t>& coor)
{
  if (shape.size() != coor.size()) {
    throw Error(__FILE__, __LINE__, "get_direct_index: shape and coor dimentionality mismatch");
  }
  for (size_t i_dim = 0; i_dim < shape.size(); i_dim++) {
    if (coor[i_dim] >= shape[i_dim]) {
      ostringstream msg;
      msg << "get_direct_index: coor beyond shape at dim " << i_dim << ": ";
      print_coor(msg, coor); msg << " "; print_coor(msg, shape);
      throw Error(__FILE__, __LINE__, msg.str());
    }
  }

  return get_direct_index__(0, shape, coor, 0);
}
